from abc import ABC, abstractmethod
from typing import Dict, Any, AsyncGenerator, Optional

from utils.SendResponseDTO import SendResponseDTO
from utils.loggerUtils import LoggerUtils
from .baseAgent import BaseAgent

logger = LoggerUtils.setup_logger(__name__)


class StreamableAgent(BaseAgent, ABC):
    """可流式处理的Agent基类，提供通用的流式处理功能"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    async def _yield_thinking_message(self, message: str, stage: str = "thinking") -> Dict[str, Any]:
        """发送思考过程消息的通用方法"""
        return SendResponseDTO.create_thinking_message(message, stage)

    async def _yield_error_message(self, error_msg: str, stage: str = "processing") -> Dict[str, Any]:
        """发送错误消息的通用方法"""
        LoggerUtils.log_error(error_msg, logger=logger)
        return SendResponseDTO.create_error_message(error_msg, stage)

    async def _yield_final_result(self, success: bool, response: str,
                                  error: Optional[str] = None) -> Dict[str, Any]:
        """发送最终结果的通用方法"""
        LoggerUtils.log_agent_complete(self.agent_name, len(response), success, logger)
        return SendResponseDTO.create_final_result_message(success, response, self.agent_name, error)

    @abstractmethod
    async def process_stream(self, input_data: Dict[str, Any]) -> AsyncGenerator[Dict[str, Any], None]:
        """抽象的流式处理方法，子类必须实现"""
        pass

        # ==================== 通用图流式处理方法 ====================

    async def _process_graph_stream(self, graph_stream, config: Dict[str, Any]) -> Any:
        """通用的图流式处理方法 - 可被子类重写以自定义处理逻辑"""
        current_state = {}

        async for chunk in graph_stream:
            # 检查是否有中断信息
            if "__interrupt__" in chunk:
                interaction_message = self._handle_interrupt(chunk, config)
                if interaction_message:
                    yield interaction_message
                    return

            # 处理正常的节点更新
            for node_name, node_result in chunk.items():
                logger.info(f"🔄 执行节点: {node_name}")
                current_state.update(node_result)

                # finish 节点不发送状态消息 - 子类可以重写此逻辑
                if self._should_skip_node_update(node_name):
                    continue

                # 发送节点更新消息 - 子类可以重写以自定义消息格式
                node_update = await self._send_node_update(node_name, node_result)
                if node_update:
                    yield node_update

        # 流式执行完成，检查最终状态
        final_response = current_state.get("final_response", "")
        success = current_state.get("success", False)

        yield await self._yield_final_result(
            success=success,
            response=final_response
        )

    def _should_skip_node_update(self, node_name: str) -> bool:
        """判断是否跳过节点更新消息 - 基础版本只跳过finish节点"""
        return "finish" in node_name.lower()

    async def _send_node_update(self, node_name: str, node_result: Dict[str, Any]) -> Dict[str, Any]:
        """发送节点更新消息 - 基础版本提供通用格式"""
        from datetime import datetime

        # 默认的通用节点更新消息
        node_messages = {
            "think": "🤔 正在思考...",
            "execute_tools": "🔧 正在调用工具...",
            "execute_mcp": "🔌 正在调用MCP服务...",
            "execute_rag": "📚 正在调用RAG检索...",
        }

        base_message = node_messages.get(node_name, f"正在执行 {node_name}...")

        # 对于思考节点，尝试使用具体的思考内容
        if node_name == "think" and node_result:
            think_content = node_result.get("current_thought", "")
            if think_content:
                base_message = think_content

        # 创建通用的节点更新消息
        return {
            "type": "node_update",
            "content": base_message,
            "timestamp": datetime.now().isoformat(),
            "data": {
                "node_name": node_name,
                "stage": f"node_{node_name}"
            }
        }

    def _handle_interrupt(self, chunk: Dict[str, Any], config: Dict[str, Any]) -> Dict[str, Any]:
        """处理中断信息 - 基础版本，子类可以重写以提供具体实现"""
        logger.warning("⚠️ 基础Agent不支持中断处理，请在子类中重写此方法")
        return None
